{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cc126ad2",
   "metadata": {},
   "source": [
    "# 数据准备"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d176ab4c",
   "metadata": {},
   "source": [
    "# !wget https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f3a185f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "文本长度: 1115394\n",
      "文本前100个字符:\n",
      "First Citizen:\n",
      "Before we proceed any further, hear me speak.\n",
      "\n",
      "All:\n",
      "Speak, speak.\n",
      "\n",
      "First Citizen:\n",
      "You\n",
      "字典大小: 65\n",
      "字典内容: ['\\n', ' ', '!', '$', '&', \"'\", ',', '-', '.', '3', ':', ';', '?', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']\n",
      "\n",
      "字符到索引的映射示例:\n",
      "'F' -> 18\n",
      "'i' -> 47\n",
      "'r' -> 56\n",
      "'s' -> 57\n",
      "'t' -> 58\n",
      "' ' -> 1\n",
      "'C' -> 15\n",
      "'i' -> 47\n",
      "'t' -> 58\n",
      "'i' -> 47\n",
      "'z' -> 64\n",
      "'e' -> 43\n",
      "'n' -> 52\n",
      "':' -> 10\n",
      "'\n",
      "' -> 0\n",
      "'B' -> 14\n",
      "'e' -> 43\n",
      "'f' -> 44\n",
      "'o' -> 53\n",
      "'r' -> 56\n",
      "\n",
      "文本转换为数字序列的前20个元素:\n",
      "[18 47 56 57 58  1 15 47 58 47 64 43 52 10  0 14 43 44 53 56]\n",
      "将数字序列转回字符:\n",
      "First Citizen:\n",
      "Befor\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "# 读取Shakespeare文本文件\n",
    "with open('shakespeare.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()\n",
    "\n",
    "# 打印文本的前100个字符\n",
    "print(f\"文本长度: {len(text)}\")\n",
    "print(f\"文本前100个字符:\\n{text[:100]}\")\n",
    "\n",
    "# 创建字符级别的字典\n",
    "vocab = sorted(set(text))\n",
    "print(f\"字典大小: {len(vocab)}\")\n",
    "print(f\"字典内容: {vocab}\")\n",
    "\n",
    "# 创建字符到索引的映射\n",
    "char_to_idx = {char: idx for idx, char in enumerate(vocab)}\n",
    "idx_to_char = {idx: char for idx, char in enumerate(vocab)}\n",
    "\n",
    "# 打印映射示例\n",
    "print(\"\\n字符到索引的映射示例:\")\n",
    "for char in text[:20]:\n",
    "    print(f\"'{char}' -> {char_to_idx[char]}\")\n",
    "\n",
    "# 将文本转换为数字序列\n",
    "text_as_int = np.array([char_to_idx[c] for c in text]) #把全部文本都变为id\n",
    "print(f\"\\n文本转换为数字序列的前20个元素:\\n{text_as_int[:20]}\")\n",
    "print(f\"将数字序列转回字符:\\n{''.join([idx_to_char[idx] for idx in text_as_int[:20]])}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a08bc2a",
   "metadata": {},
   "source": [
    "# 把莎士比亚文集分成一个一个的样本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e2dd4a94",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "输入批次形状: torch.Size([64, 100])\n",
      "目标批次形状: torch.Size([64, 100])\n",
      "tensor([[21, 15, 20,  ...,  1, 51, 53],\n",
      "        [56, 39, 52,  ..., 53, 58,  1],\n",
      "        [63, 53, 59,  ..., 52, 42,  1],\n",
      "        ...,\n",
      "        [ 1, 46, 39,  ...,  1, 50, 53],\n",
      "        [17,  1, 27,  ..., 18, 21, 32],\n",
      "        [47, 52,  1,  ..., 39, 42,  1]], dtype=torch.int32)\n",
      "tensor([[15, 20, 13,  ..., 51, 53, 56],\n",
      "        [39, 52, 45,  ..., 58,  1, 57],\n",
      "        [53, 59, 56,  ..., 42,  1, 40],\n",
      "        ...,\n",
      "        [46, 39, 58,  ..., 50, 53, 53],\n",
      "        [ 1, 27, 18,  ..., 21, 32, 38],\n",
      "        [52,  1, 58,  ..., 42,  1, 45]], dtype=torch.int32)\n",
      "\n",
      "数据集大小: 11043\n",
      "批次数量: 172\n"
     ]
    }
   ],
   "source": [
    "# 定义序列长度和批次大小\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "seq_length = 100  # 每个样本的序列长度\n",
    "batch_size = 64   # 每个批次的样本数量\n",
    "\n",
    "# 创建自定义数据集类\n",
    "class ShakespeareDataset(Dataset):\n",
    "    def __init__(self, text_as_int, seq_length):\n",
    "        self.text_as_int = text_as_int\n",
    "        self.seq_length = seq_length\n",
    "        self.sub_len = seq_length + 1 #一个样本的长度\n",
    "        \n",
    "    def __len__(self):\n",
    "        # 计算可能的序列数量\n",
    "        return len(self.text_as_int)//(self.seq_length+1) #+1是因为要预测下一个字符\n",
    "        \n",
    "    def __getitem__(self, idx):\n",
    "        return self.text_as_int[idx*self.sub_len:(idx+1)*self.sub_len]\n",
    "\n",
    "# 定义collate函数，用于处理批次数据\n",
    "def collate_fct(batch):\n",
    "    # 将批次数据转换为张量\n",
    "    batch = torch.tensor(batch)\n",
    "    # 输入序列是除了最后一个字符的所有字符\n",
    "    input_batch = batch[:, :-1]\n",
    "    # 目标序列是除了第一个字符的所有字符\n",
    "    target_batch = batch[:, 1:]\n",
    "    return input_batch, target_batch\n",
    "\n",
    "# 创建数据集实例\n",
    "shakespeare_dataset = ShakespeareDataset(text_as_int, seq_length)\n",
    "\n",
    "# 创建数据加载器\n",
    "dataloader = DataLoader(shakespeare_dataset, batch_size=batch_size, shuffle=True, drop_last=True, collate_fn=collate_fct)\n",
    "\n",
    "# 打印示例，查看输入和目标\n",
    "for input_batch, target_batch in dataloader:\n",
    "    print(f\"输入批次形状: {input_batch.shape}\")\n",
    "    print(f\"目标批次形状: {target_batch.shape}\")\n",
    "    \n",
    "    # 打印第一个样本的输入和目标\n",
    "    print(input_batch)\n",
    "    print(target_batch)\n",
    "    break\n",
    "\n",
    "print(f\"\\n数据集大小: {len(shakespeare_dataset)}\")\n",
    "print(f\"批次数量: {len(dataloader)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0606fa96",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "172"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "11043//64"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0353bb1a",
   "metadata": {},
   "source": [
    "# 搭建模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "50a9336b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ShakespeareRNN(\n",
      "  (embedding): Embedding(65, 256)\n",
      "  (rnn): RNN(256, 1024, batch_first=True)\n",
      "  (dense): Linear(in_features=1024, out_features=65, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# 定义RNN模型\n",
    "class ShakespeareRNN(nn.Module):\n",
    "    def __init__(self, vocab_size, embedding_dim, hidden_dim, batch_size):\n",
    "        super(ShakespeareRNN, self).__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, embedding_dim)\n",
    "        self.rnn = nn.RNN(\n",
    "            embedding_dim,\n",
    "            hidden_dim,\n",
    "            num_layers=1,\n",
    "            bidirectional=False,\n",
    "            batch_first=True\n",
    "        )\n",
    "        self.dense = nn.Linear(hidden_dim, vocab_size)\n",
    "        \n",
    "    def forward(self, x, hidden=None):\n",
    "        # 输入形状: [batch_size, sequence_length]\n",
    "        x = self.embedding(x)  # 形状: [batch_size, sequence_length, embedding_dim]\n",
    "        output, hidden = self.rnn(x, hidden)  # 形状: [batch_size, sequence_length, hidden_dim]\n",
    "        output = self.dense(output)  # 形状: [batch_size, sequence_length, vocab_size]\n",
    "        return output, hidden\n",
    "    \n",
    "\n",
    "\n",
    "# 定义模型参数\n",
    "vocab_size = len(char_to_idx)  # 词汇表大小\n",
    "embedding_dim = 256  # 嵌入维度\n",
    "rnn_units = 1024  # RNN单元数量\n",
    "\n",
    "# 实例化模型\n",
    "model = ShakespeareRNN(vocab_size, embedding_dim, rnn_units, batch_size)\n",
    "print(model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8ef60cf5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "输入形状: torch.Size([4, 100])\n",
      "输出形状: torch.Size([4, 100, 65])\n",
      "模型前向计算验证成功！\n"
     ]
    }
   ],
   "source": [
    "# 创建一个小批量数据来测试模型\n",
    "batch_size = 4\n",
    "seq_length = 100\n",
    "test_input = torch.randint(0, vocab_size, (batch_size, seq_length))\n",
    "\n",
    "# 进行前向计算\n",
    "with torch.no_grad():\n",
    "    output, hidden = model(test_input)\n",
    "    \n",
    "# 打印输出形状\n",
    "print(f\"输入形状: {test_input.shape}\")\n",
    "print(f\"输出形状: {output.shape}\")\n",
    "\n",
    "# 验证输出是否符合预期\n",
    "assert output.shape == (batch_size, seq_length, vocab_size), \"输出形状不符合预期\"\n",
    "assert hidden.shape == (1, batch_size, rnn_units), \"隐藏状态形状不符合预期\"\n",
    "\n",
    "print(\"模型前向计算验证成功！\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95e6a3f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义损失函数和优化器\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# 训练函数\n",
    "def train_step(model, input_batch, target_batch, optimizer, criterion):\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    # 前向传播\n",
    "    output, _ = model(input_batch)\n",
    "    \n",
    "    # 计算损失\n",
    "    # 重塑输出和目标以适应CrossEntropyLoss\n",
    "    output = output.reshape(-1, vocab_size)\n",
    "    target_batch = target_batch.reshape(-1)\n",
    "    \n",
    "    loss = criterion(output, target_batch)\n",
    "    \n",
    "    # 反向传播\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    return loss.item()\n",
    "\n",
    "# 生成文本函数\n",
    "def generate_text(model, start_string, char2idx, idx2char, num_generate=1000, temperature=1.0):\n",
    "    model.eval()\n",
    "    \n",
    "    # 将起始字符串转换为数字\n",
    "    input_eval = [char2idx[s] for s in start_string]\n",
    "    input_eval = torch.tensor([input_eval], dtype=torch.long)\n",
    "    \n",
    "    # 空字符串用于存储结果\n",
    "    text_generated = []\n",
    "    \n",
    "    # 隐藏状态形状: (1, batch_size, rnn_units)\n",
    "    hidden = None\n",
    "    \n",
    "    for i in range(num_generate):\n",
    "        # 前向传播\n",
    "        predictions, hidden = model(input_eval, hidden)\n",
    "        \n",
    "        # 使用最后一个时间步的预测\n",
    "        predictions = predictions[:, -1, :] / temperature\n",
    "        \n",
    "        # 使用分类分布采样预测的字符\n",
    "        probabilities = F.softmax(predictions, dim=-1)\n",
    "        predicted_id = torch.multinomial(probabilities, num_samples=1)\n",
    "        \n",
    "        # 将预测的字符作为下一个输入\n",
    "        input_eval = predicted_id\n",
    "        \n",
    "        # 添加预测的字符到生成的文本\n",
    "        text_generated.append(idx2char[predicted_id.item()])\n",
    "    \n",
    "    return start_string + ''.join(text_generated)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
